#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import cv2
from PIL import Image
import numpy as np
import os
import glob
import argparse
from pathlib import Path


class VideoFrameExtractor:
    """
    VideoFrameExtractor
    ==================

    功能：
        从视频文件中按固定间隔抽取帧，并保存为图像文件。

    输入：
        - video_folder: 视频文件夹路径
        - output_base_dir: 输出目录
        - frame_interval: 抽帧间隔

    输出：
        - 按视频文件名组织的图像文件夹

    使用方法：
        extractor = VideoFrameExtractor(video_folder, output_base_dir, frame_interval)
        extractor()
    """

    def __init__(self, video_folder, output_base_dir, frame_interval=600):
        self.video_folder = Path(video_folder)
        self.output_base_dir = Path(output_base_dir)
        self.frame_interval = frame_interval

        # 确保输出基目录存在
        os.makedirs(self.output_base_dir, exist_ok=True)

        # 支持的视频文件扩展名
        self.video_extensions = ['*.avi', '*.mp4', '*.mov', '*.mkv', '*.flv', '*.wmv']

    def _get_video_files(self):
        """获取所有视频文件"""
        video_files = []
        for extension in self.video_extensions:
            video_files.extend(glob.glob(str(self.video_folder / extension)))
        return video_files

    def _extract_frames_from_video(self, video_path):
        """从单个视频文件中提取帧"""
        cap = cv2.VideoCapture(video_path)
        isOpened = cap.isOpened()

        if not isOpened:
            print(f"无法打开视频: {video_path}")
            return 0

        # 从视频路径中提取基础文件名（不含扩展名）
        base_filename = Path(video_path).stem

        # 为当前视频创建专属输出目录
        output_dir = self.output_base_dir / base_filename
        os.makedirs(output_dir, exist_ok=True)

        imageNum = 0
        frame_count = 0

        print(f"开始处理视频: {video_path}")

        while isOpened:
            (frameState, frame) = cap.read()

            if not frameState:
                break

            if frame_count % self.frame_interval == 0:
                # 格式转变，BGRtoRGB
                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                # 转变成Image
                frame = Image.fromarray(np.uint8(frame))
                frame = np.array(frame)
                # RGBtoBGR满足opencv显示格式
                frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)

                # 使用原始文件名前缀 + 帧数作为图片名
                fileName = output_dir / f'{base_filename}_{frame_count:06d}.jpg'
                cv2.imwrite(str(fileName), frame, [cv2.IMWRITE_JPEG_QUALITY, 100])
                imageNum += 1

            frame_count += 1

        cap.release()
        print(f"完成视频 {video_path} 的处理，共保存 {imageNum} 张图片")
        return imageNum

    def extract_frames(self):
        """提取所有视频的帧"""
        video_files = self._get_video_files()
        print(f"找到 {len(video_files)} 个视频文件")

        total_images = 0

        # 处理每个视频文件
        for video_path in video_files:
            images_count = self._extract_frames_from_video(video_path)
            total_images += images_count

        print(f'所有视频处理完成，共提取 {total_images} 张图片!')

    def __call__(self):
        self.extract_frames()


def main():
    parser = argparse.ArgumentParser(description='视频抽帧工具')
    parser.add_argument('video_folder', type=str, help='视频文件夹路径')
    parser.add_argument('output_base_dir', type=str, help='输出目录')
    parser.add_argument('--frame_interval', type=int, default=600,
                        help='抽帧间隔（默认：600帧）')

    args = parser.parse_args()

    extractor = VideoFrameExtractor(
        video_folder=args.video_folder,
        output_base_dir=args.output_base_dir,
        frame_interval=args.frame_interval
    )
    extractor()


if __name__ == "__main__":
    main()